"""
Keyboard input controller for Isaac Sim.
"""

import numpy as np
import carb
import omni


class DirectKeyboardController:
    """直接使用Isaac Sim键盘事件系统"""

    def __init__(self):
        # 6维命令数组：[I, K, J, L, U, O] 对应原始InternUtopia实现
        self.command = np.array([0.0, 0.0, 0.0, 0.0, 0.0, 0.0])
        self.subscribe()

    def subscribe(self):
        """订阅Isaac Sim键盘事件"""
        app_window = omni.appwindow.get_default_app_window()
        key_input = carb.input.acquire_input_interface()
        key_input.subscribe_to_keyboard_events(app_window.get_keyboard(), self._on_keyboard_event)
        print("键盘事件监听已启用: I/K/J/L/U/O")

    def _on_keyboard_event(self, event, *args, **kwargs):
        """处理键盘事件 - 完全复制原始实现"""
        if (event.type == carb.input.KeyboardEventType.KEY_PRESS or
            event.type == carb.input.KeyboardEventType.KEY_REPEAT):

            if event.input == carb.input.KeyboardInput.I:
                self.command = np.array([1.0, 0.0, 0.0, 0.0, 0.0, 0.0])
            elif event.input == carb.input.KeyboardInput.K:
                self.command = np.array([0.0, 1.0, 0.0, 0.0, 0.0, 0.0])
            elif event.input == carb.input.KeyboardInput.J:
                self.command = np.array([0.0, 0.0, 1.0, 0.0, 0.0, 0.0])
            elif event.input == carb.input.KeyboardInput.L:
                self.command = np.array([0.0, 0.0, 0.0, 1.0, 0.0, 0.0])
            elif event.input == carb.input.KeyboardInput.U:
                self.command = np.array([0.0, 0.0, 0.0, 0.0, 1.0, 0.0])
            elif event.input == carb.input.KeyboardInput.O:
                self.command = np.array([0.0, 0.0, 0.0, 0.0, 0.0, 1.0])

        elif event.type == carb.input.KeyboardEventType.KEY_RELEASE:
            self.command = np.array([0.0, 0.0, 0.0, 0.0, 0.0, 0.0])

        return True

    def get_speed_commands(self):
        """计算速度命令 - I-K, J-L, U-O"""
        x_speed = self.command[0] - self.command[1]  # I - K: 前进/后退
        y_speed = self.command[2] - self.command[3]  # J - L: 左移/右移
        z_speed = self.command[4] - self.command[5]  # U - O: 左转/右转
        return x_speed, y_speed, z_speed